import warnings, argparse, torch
warnings.filterwarnings('ignore')
import numpy as np
import torch.nn as nn
import matplotlib.pylab as plt
from models.common import DetectMultiBackend

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--base-weights', type=str, default=r'sumdata_EfficientViT_b1_dyhead2\weights\best.pt', help='base weights path')
    parser.add_argument('--prune-weights', type=str, default=r'yolov5s_EfficientViT_dyhead_group_taylor_exp12\weights\prune.pt', help='prune weights path')
    opt = parser.parse_args()
    
    print(f'loading base model from {opt.base_weights}')
    base_model = DetectMultiBackend(opt.base_weights, device=torch.device('cpu'))
    print(f'loading prune model from {opt.prune_weights}')
    prune_model = DetectMultiBackend(opt.prune_weights, device=torch.device('cpu'))
    
    base_model_state_dict, prune_model_state_dict = base_model.state_dict(), prune_model.state_dict()
    
    base_channels, prune_channels, channels_dis, names = [], [], [], []
    for (base_m_name, base_m), (prune_m_name, prune_m) in zip(base_model.named_modules(), prune_model.named_modules()):
        try:
            if isinstance(base_m, nn.Conv2d):
                names.append(base_m_name.replace('model.model.', ''))
                base_channels.append(base_m.out_channels)
                prune_channels.append(prune_m.out_channels)
                channels_dis.append(base_m.out_channels - prune_m.out_channels)
            elif isinstance(base_m, nn.Linear):
                names.append(base_m_name.replace('model.model.', ''))
                base_channels.append(base_m.out_features)
                prune_channels.append(prune_m.out_features)
                channels_dis.append(base_m.out_features - prune_m.out_features)
        except:
            continue
    def simplify_layer_name(name):
        if name == '0.input_stem.op_list.0.conv':
            parts = name.split('.')
            return parts[-1]
        else:
            # 对以 '0.' 开头的层进行特殊处理
            if name.startswith('0.'):
                parts = name.split('.')
                # 仅保留关键部分
                if 'conv' in parts:
                    return '.'.join(parts[-2:])  # 例如 'input_stem.op_list'
                elif 'aggreg' in parts:
                    return '.'.join(parts[-3:])  # 例如 'stages.0.op_list'
                
                # elif 'dyhead' in parts:
                #     return '.'.join(parts[1:4])  # 例如 'dyhead.0.spatial_conv_offset'
            elif name.startswith('19.'):
                parts = name.split('.')
                if 'spatial_conv_offset' in parts:
                    return parts[-1] 
                elif 'scale_attn_module' in parts:
                    return '.'.join(parts[-2:])
                elif 'task_attn_module' in parts:
                    return '.'.join(parts[-3:]) 
        return name  # 其他层保持原样或进行适当简化         
    
    base_channels, prune_channels, channels_dis, names = np.array(base_channels), np.array(prune_channels), np.array(channels_dis), np.array(names)
    channels_dis_sort = np.argsort(channels_dis)[::-1]
    

    filtered_base_channels = []
    filtered_prune_channels = []
    filtered_names = []

    for (name, base_ch, prune_ch) in zip(names, base_channels, prune_channels):
        # print(name)
            simplified_name = simplify_layer_name(name)
            if simplified_name:  # 只有当层名称被保留时，才收集通道数
                filtered_names.append(simplified_name)
                filtered_base_channels.append(base_ch)
                filtered_prune_channels.append(prune_ch)
    # print(f'filtered_names:{filtered_names}\n')
    # print(f'filtered_base_channels:{filtered_base_channels}\n')
    # print(f'filtered_prune_channels:{filtered_prune_channels}\n')
    
    
    x = np.arange(len(filtered_names))
    plt.figure(figsize=(35, 15))
    plt.xticks(x, filtered_names, rotation=90,fontsize=12)
    plt.bar(x, filtered_base_channels, color='skyblue', label='Before pruning')
    plt.bar(x, filtered_prune_channels, color='lightgreen', label='After pruning')
    plt.legend(fontsize=30)
    plt.title('Comparison before and after pruning', fontdict={'fontsize':30})
    plt.tight_layout()
    for i, (base_ch, prune_ch) in enumerate(zip(base_channels, prune_channels)):
        plt.text(i, base_ch, str(base_ch), ha='center', va='bottom', rotation=90)
        plt.text(i, prune_ch, str(prune_ch), ha='center', va='bottom', rotation=90)

    plt.savefig('channels_chart', dpi = 300)
    print('save to channels_chart.png')


#--------------------排序后通道图------------------------    
    # base_channels, prune_channels, names = base_channels[channels_dis_sort], prune_channels[channels_dis_sort], names[channels_dis_sort]


    # filtered_base_channels = []
    # filtered_prune_channels = []
    # filtered_names = []

    # for (name, base_ch, prune_ch) in zip(names, base_channels, prune_channels):
    #     # print(name)
    #         simplified_name = simplify_layer_name(name)
    #         if simplified_name:  # 只有当层名称被保留时，才收集通道数
    #             filtered_names.append(simplified_name)
    #             filtered_base_channels.append(base_ch)
    #             filtered_prune_channels.append(prune_ch)


    # x = np.arange(len(filtered_names))
    # plt.figure(figsize=(25, 8))
    # plt.xticks(x, filtered_names, rotation=90)
    # plt.bar(x, filtered_base_channels, color='skyblue', label='base')
    # plt.bar(x, filtered_prune_channels, color='lightgreen', label='prune')
    # plt.legend(fontsize=20)
    # plt.title('Channel contrast diagram', fontdict={'fontsize':20})
    # plt.tight_layout()
    # plt.savefig('channels_chart_sort')
    # print('save to channels_chart_sort.png')